from typing import List, Dict
import os
import pandas as pd
from datetime import datetime
from vnpy.trader.object import BarData
from vnpy.trader.constant import Exchange, Interval

import os

# 获取当前脚本所在目录
current_dir = os.path.dirname(os.path.abspath(__file__))
# 获取项目根目录
project_root = os.path.dirname(os.path.dirname(current_dir))


class DataStorage:
    """数据存储类"""
    
    def __init__(self):
        """初始化"""
        # 获取项目根目录
        self.data_path = os.path.join(project_root, "portfolio_trader\data")
        os.makedirs(self.data_path, exist_ok=True)

    def _get_exchange(self, symbol: str) -> Exchange:
        """获取交易所"""
        if symbol.startswith(("600", "601", "603", "688")):
            return Exchange.SSE
        elif symbol.startswith(("000", "002", "300")):
            return Exchange.SZSE
        elif symbol.startswith(("4", "8")):  # 北交所
            return Exchange.BSE
        else:
            return Exchange.SSE  # 默认返回上交所 
    
    def has_data(self, symbol: str, data_type: str = "daily") -> bool:
        """检查是否有指定股票的数据"""
        try:
            # 检查CSV文件是否存在
            file_path = os.path.join(self.data_path, f"{symbol}_{data_type}.csv")
            if os.path.exists(file_path):
                # 检查文件是否为空
                df = pd.read_csv(file_path)
                return not df.empty
            return False
        except Exception as e:
            print(f"检查数据失败: {str(e)}")
            return False

    def _save_new_csv_bar(self, bars: List[BarData], csv_path: str) -> None:
        """保存新数据到CSV文件"""
        # 转换为DataFrame
        data = []
        for bar in bars:
            data.append({
                "symbol": bar.symbol,
                "exchange": bar.exchange.value,
                "datetime": bar.datetime.replace(hour=15),
                "interval": bar.interval.value,
                "volume": bar.volume,
                "turnover": bar.turnover,
                "open_interest": bar.open_interest,
                "open_price": bar.open_price,
                "high_price": bar.high_price,
                "low_price": bar.low_price,
                "close_price": bar.close_price,
                "vt_symbol": f"{bar.symbol}.{bar.exchange.value}"
            })
        
        df = pd.DataFrame(data).sort_values("datetime", ascending=False)
        
        # 保存到CSV
        try:
            df.to_csv(csv_path, index=False)
            print(f"成功保存 {len(df)} 条数据到 {csv_path}")
        except Exception as e:
            print(f"保存数据到CSV文件失败: {str(e)}") 

    def save_bar_data(self, bars: List[BarData], symbol: str) -> None:
        """保存K线数据"""
        if not bars:
            print("警告: 没有数据可保存")
            return
        
        csv_path = os.path.join(self.data_path, f"{symbol}_daily.csv")
        
        # 情况一：本地无文件时直接保存
        if not os.path.exists(csv_path):
            self._save_new_csv_bar(bars, csv_path)
            return
        
        # 情况二：本地有文件时进行合并操作
        try:
            # 读取本地历史数据
            existing_df = pd.read_csv(csv_path)
            existing_df["datetime"] = pd.to_datetime(existing_df["datetime"])
            
            # 转换新数据为DataFrame
            new_data = []
            for bar in bars:
                new_data.append({
                    "symbol": bar.symbol,
                    "exchange": bar.exchange.value,
                    "datetime": bar.datetime,
                    "interval": bar.interval.value,
                    "volume": bar.volume,
                    "turnover": bar.turnover,
                    "open_interest": bar.open_interest,
                    "open_price": bar.open_price,
                    "high_price": bar.high_price,
                    "low_price": bar.low_price,
                    "close_price": bar.close_price,
                    "vt_symbol": f"{bar.symbol}.{bar.exchange.value}"
                })
            new_df = pd.DataFrame(new_data)
            new_df["datetime"] = pd.to_datetime(new_df["datetime"])
            
            # 合并数据（保留最新记录）并按时间降序排序
            combined_df = pd.concat([existing_df, new_df]).drop_duplicates(
                subset=["datetime"], keep="last"
            ).sort_values("datetime", ascending=False)
            
            # 保存合并后的数据
            combined_df.to_csv(csv_path, index=False)
            print(f"成功合并保存 {len(combined_df)} 条数据到 {csv_path}")
                
        except Exception as e:
            print(f"保存K线数据时发生错误: {str(e)}")

    def save_valuation_data(self, valuation: pd.DataFrame, symbol: str) -> None:
        """保存估值数据"""
        if valuation.empty:
            print("警告: 没有估值数据可保存")
            return
        
        csv_path = os.path.join(self.data_path, f"{symbol}_valuation.csv")
        
        # 情况一：本地无文件时直接保存
        if not os.path.exists(csv_path):
            print("文件不存在，直接保存估值数据")
            try:
                # 按时间倒序排序后保存
                valuation = valuation.sort_values("datetime", ascending=False)
                valuation.to_csv(csv_path, index=False)
                print(f"成功保存估值数据到 {csv_path}")
            except Exception as e:
                print(f"保存估值数据到CSV文件失败: {str(e)}")
            return
        
        # 情况二：本地有文件时进行合并操作
        try:
            # 读取本地历史数据
            existing_df = pd.read_csv(csv_path)
            existing_df["datetime"] = pd.to_datetime(existing_df["datetime"])
            # 转换时间格式
            valuation["datetime"] = pd.to_datetime(valuation["datetime"])
            
            # 合并数据（保留最新记录）并按时间倒序排序
            combined_df = pd.concat([existing_df, valuation]).drop_duplicates(
                subset=["datetime"], keep="last"
            ).sort_values("datetime", ascending=False)
            
            # 保存合并后的数据
            combined_df.to_csv(csv_path, index=False)
            print(f"成功合并保存估值数据到 {csv_path}")
                
        except Exception as e:
            print(f"保存估值数据时发生错误: {str(e)}")
